import numpy as np
from math import sin, cos, radians, sqrt
from typing import Tuple

# --- I. 物理常数和弹丸参数 ---
RHO0 = 1.225         # 海平面空气密度 (kg/m^3)
T0 = 288.15          # 海平面温度 (K)
L = 0.0065           # 绝热垂直温度梯度 (K/m)
G = 9.80665          # 重力加速度 (m/s^2)
R = 287.05           # 干燥空气比气体常数 (J/kg·K)
P_TERM = (G / (R * L)) - 1 # 用于计算密度比的指数
C_SOUND = 338.0      # 海平面声速 (m/s)

M = 5.0              # 弹丸质量 (kg)
A = 0.0095           # 弹丸截面积 (m^2)

MA_POINTS = np.array([0.0, 0.7, 1.0, 1.3, 2.0, 2.5]) # 马赫数查找点
CD_POINTS = np.array([0.48, 0.45, 0.95, 0.45, 0.40, 0.40]) # 阻力系数

# --- II. 动态物理模型和RK4函数 (2D简化版) ---

def get_drag_coefficient(v_rel):
    """根据相对速度计算空气阻力系数 (Cd)。"""
    mach_number = v_rel / C_SOUND
    return np.interp(mach_number, MA_POINTS, CD_POINTS)

def get_air_density(z):
    """根据高度 Z 计算空气密度。"""
    h = max(0.0, z) 
    temp_ratio = (T0 - L * h) / T0
    if temp_ratio <= 0: return 0.0 
    rho = RHO0 * (temp_ratio) ** P_TERM
    return rho

def derivatives_of_state_2d(S: np.ndarray, v_wind_x: float, params: Tuple[float, float, float]) -> np.ndarray:
    """计算 2D 状态向量 S = [x, z, vx, vz] 在时间上的微分。"""
    x, z, vx, vz = S
    M, A, G = params 
    
    # 1. 计算相对速度 (相对空气)
    vx_rel = vx - v_wind_x
    vz_rel = vz
    v_rel = sqrt(vx_rel**2 + vz_rel**2) # 相对速度的模

    if v_rel < 1e-6:
        return np.array([vx, vz, 0.0, -G])

    rho = get_air_density(z)
    Cd = get_drag_coefficient(v_rel)

    # 2. 空气阻力加速度
    K = -(Cd * rho * A * v_rel) / (2.0 * M)
    
    dvx_drag = K * vx_rel
    dvz_drag = K * vz_rel

    # 3. 合并加速度
    dvx_dt = dvx_drag
    dvz_dt = dvz_drag - G

    return np.array([vx, vz, dvx_dt, dvz_dt])

def rk4_step_2d(S: np.ndarray, dt: float, vwx: float, params: Tuple[float, float, float]) -> np.ndarray:
    """四阶龙格-库塔方法进行一步 2D 积分。"""
    f = lambda S: derivatives_of_state_2d(S, vwx, params)
    
    K1 = dt * f(S)
    K2 = dt * f(S + K1 / 2.0)
    K3 = dt * f(S + K2 / 2.0)
    K4 = dt * f(S + K3)
    
    S_new = S + (K1 + 2.0 * K2 + 2.0 * K3 + K4) / 6.0
    return S_new

def simulate_and_get_impact_2d(v0: float, theta_deg: float, dt: float, v_wind_x: float) -> float:
    """模拟 2D 轨迹并返回最终落点 X 坐标 (射程)。"""
    params = (M, A, G)
    theta_rad = radians(theta_deg)
    
    # 2D 状态向量: [x, z, vx, vz]
    S = np.array([
        0.0, 0.0, 
        v0 * cos(theta_rad), 
        v0 * sin(theta_rad)
    ])
    
    while S[1] >= -0.01:
        S_new = rk4_step_2d(S, dt, v_wind_x, params) 
        
        if S_new[1] < 0 and S[1] >= 0:
            t_impact = S[1] / (S[1] - S_new[1])
            R_final = S[0] + (S_new[0] - S[0]) * t_impact
            return R_final
            
        S = S_new
        
    return S[0]

# --- III. 搜索函数 (计算 V0 和 theta) ---

def binary_search_v0_for_range(target_range: float, theta_deg: float, dt: float, v_wind_x: float, V_low=10.0, V_high=2000.0, tolerance=1.0) -> float:
    """查找固定角度和目标距离所需的 V0 (考虑纵向风 v_wind_x)。"""
    
    for _ in range(50):
        V_mid = (V_low + V_high) / 2.0
        R_mid = simulate_and_get_impact_2d(V_mid, theta_deg, dt, v_wind_x)
        
        if abs(R_mid - target_range) < tolerance:
            return V_mid
        
        if R_mid < target_range:
            V_low = V_mid 
        else:
            V_high = V_mid 
            
    # 如果搜索失败，返回一个指示失败的值，而不是抛出异常，以便主程序可以处理
    return np.nan

def search_min_v0_and_theta(target_range: float, dt: float, v_wind_x: float) -> Tuple[float, float]:
    """
    搜索达到目标射程所需的最小初速度 V0 及其对应的发射角 theta。
    """
    
    # 角度搜索范围和步长
    THETA_MIN = 15.0
    THETA_MAX = 60.0
    THETA_STEP = 0.5  # 角度步长，提供较好的精度

    best_v0 = float('inf')
    best_theta = np.nan
    
    # 遍历可能的发射角
    for theta_deg in np.arange(THETA_MIN, THETA_MAX + THETA_STEP, THETA_STEP):
        # 搜索当前角度所需的 V0
        v0_required = binary_search_v0_for_range(
            target_range=target_range, 
            theta_deg=theta_deg, 
            dt=dt, 
            v_wind_x=v_wind_x
        )
        
        # 检查是否找到有效的 V0 并且比当前最优 V0 小
        if not np.isnan(v0_required) and v0_required < best_v0:
            best_v0 = v0_required
            best_theta = theta_deg
                
    if best_v0 == float('inf'):
        # 此时 V0_opt, Theta_opt 仍为 np.nan
        return np.nan, np.nan
        
    return best_v0, best_theta

# --- IV. 主运行和表格输出 ---

def generate_optimized_table():
    
    # --- 固定参数 ---
    DT = 0.1             # RK4 步长 (s)

    # --- 表格参数 ---
    R_MIN = 1000.0
    R_MAX = 1500.0
    R_STEP = 100.0        # 目标距离步长
    V_W_MIN = -30.0       # 最小纵向风速度 (逆风)
    V_W_MAX = 30.0        # 最大纵向风速度 (顺风)
    V_W_STEP = 5.0        # 纵向风速度步长
    
    R_TARGETS = np.arange(R_MIN, R_MAX + R_STEP, R_STEP)
    V_W_TARGETS = np.arange(V_W_MIN, V_W_MAX + V_W_STEP, V_W_STEP)

    if len(R_TARGETS) < 2 or len(V_W_TARGETS) < 2:
        print("❌ 错误: 步长设置不合理，无法生成表格。")
        return

    # 存储结果的二维数组 (存储等效射程 R_equivalent)
    results_table = np.zeros((len(V_W_TARGETS), len(R_TARGETS)))

    print("--- 🎯 正在进行 2D 纵向风最优参数搜索并计算等效射程 ---")
            
    # 填充表格 (计算最优 (V0, Theta) 组合)
    for j, R_target in enumerate(R_TARGETS):
        print(f"\n[R={R_target:.0f}m] 正在搜索最优参数...")
        
        for i, V_w in enumerate(V_W_TARGETS):
            
            # 1. 搜索最优的 (V0, Theta) 组合，使其在 V_w 下射程达到 R_target
            V0_opt, Theta_opt = search_min_v0_and_theta(R_target, DT, v_wind_x=V_w)
            
            R_equivalent = np.nan
            if not np.isnan(V0_opt):
                # 2. 计算等效射程 (Equivalent Range)
                # 使用找到的最优 V0 和 Theta，但在无风 (v_wind_x=0.0) 条件下进行模拟
                try:
                    R_equivalent = simulate_and_get_impact_2d(V0_opt, Theta_opt, DT, v_wind_x=0.0)
                except Exception:
                    R_equivalent = np.nan
            
            # 3. 存储等效射程
            results_table[i, j] = R_equivalent
            
            # 打印调试信息
            V0_str = f"{V0_opt:+6.2f}" if not np.isnan(V0_opt) else "N/A"
            Theta_str = f"{Theta_opt:.1f}" if not np.isnan(Theta_opt) else "N/A"
            Req_str = f"{R_equivalent:.1f}" if not np.isnan(R_equivalent) else "N/A"
            
            print(f"  R={R_target:.0f}m, Vw={V_w:+5.1f} m/s => V₀:{V0_str}m/s, θ:{Theta_str}° | R_eq:{Req_str}m")


    # --- 结果表格输出 ---
    print("\n" + "="*95)
    print("           🎯 2D 纵向风最优发射条件下的 等效射程 ($R_{\\text{eq}}$) 查表 (单位: 米) 🎯")
    print("      [注: R_eq 是最优 (V₀, θ) 组合在无风 (V_x=0) 条件下能达到的射程]")
    print("="*95)

    # 格式化表头 (X轴：目标直线距离)
    header_cols = [f"{r:^10.0f}" for r in R_TARGETS]
    header = " | " + " | ".join(header_cols) + " |"
    print(" 纵向风 V_x (m/s) \\ 目标距离 R (m)")
    print("-" * len(header))
    print(header)
    print("-" * len(header))
    
    # 格式化数据行 (Y轴：纵向风速度)
    for i, V_w in enumerate(V_W_TARGETS):
        row = []
        for j in range(len(R_TARGETS)):
            val = results_table[i, j]
            if np.isnan(val) or val <= 0:
                row.append(f"{'N/A':^10}")
            else:
                row.append(f"{val:^10.1f}") 
        
        print(f" | {V_w:>+10.1f} | " + " | ".join(row) + " |")

    print("="*95)
    print("💡 查表指引:")
    print("   - R_eq: 为了克服纵向风 (V_x) 影响而调整 V₀ 和 θ 后，该组参数在无风时本可达到的距离。")
    print("   - R_eq > R_target: 通常发生在逆风 (V_x < 0) 情况下，需要更大的 V₀。")
    print("   - R_eq < R_target: 通常发生在顺风 (V_x > 0) 情况下，可以采用较小的 V₀。")
    
if __name__ == "__main__":
    generate_optimized_table()